#include "code_comment/code_comment_gen.h"
#include "code_comment/code_comment_info.h"

#include "clang/Basic/Diagnostic.h"
#include "clang/Basic/DiagnosticOptions.h"
#include "clang/Basic/FileManager.h"
#include "clang/Basic/IdentifierTable.h"
#include "clang/Basic/LangOptions.h"
#include "clang/Basic/SourceManager.h"
#include "clang/Basic/TokenKinds.h"
#include "clang/Frontend/TextDiagnosticPrinter.h"
#include "clang/Rewrite/Core/Rewriter.h"
#include "clang/Tooling/CommonOptionsParser.h"
#include "clang/Tooling/Tooling.h"
#include "llvm/ADT/IntrusiveRefCntPtr.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/JSON.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/ThreadPool.h"
#include "llvm/Support/raw_ostream.h"

#include <limits.h>
#include <string>
#include <vector>
#include <time.h>

using namespace llvm;
using namespace clang;

static cl::OptionCategory CodeCommentOptions("code comment common options");

static cl::opt<std::string> OutputFile("output", 
    cl::desc("output file name, it can be ignore(use default file '.code_commend_cache_file')"), 
    cl::init(".code_commend_cache_file.json"),
    cl::cat(CodeCommentOptions));
static cl::opt<unsigned> ExecConcurrency("threads",
    cl::desc("the number of threads used to process all files in parallel. Set to o for hardware concurrency"),
    cl::init(0),
    cl::Optional,
    cl::cat(CodeCommentOptions));

class CodeCommentGenRefactorFactory
{
public:
    explicit CodeCommentGenRefactorFactory(CodeCommentInfoManager& codeCommentInfoManager)
        : mCodeCommentInfoManager(codeCommentInfoManager)
        {}
    
    std::unique_ptr<clang::ASTConsumer> newASTConsumer()
    {
        return std::make_unique<CodeCommentGenASTConsumer>(mCodeCommentInfoManager);
    }

private:
    CodeCommentInfoManager& mCodeCommentInfoManager;
};

int Main_CodeCommentGen(tooling::CommonOptionsParser& OP)
{
    auto Files = OP.getSourcePathList();
    int resultMask = 0;
    std::string outFile = OutputFile.getValue();
    llvm::sys::Mutex resultSaveMutex;
    CodeCommentInfoManager codeCommentInfoManager(outFile);
    llvm::ThreadPool Pool(llvm::hardware_concurrency(ExecConcurrency));
    for(std::size_t i = 0; i < Files.size(); i++) {
        Pool.async([&, i]() {
            std::vector<std::string> subFiles = {Files[i]};
            tooling::ClangTool Tool(OP.getCompilations(), subFiles);
            CodeCommentGenRefactorFactory refactor(codeCommentInfoManager);
            std::unique_ptr<tooling::FrontendActionFactory> Factory = 
                tooling::newFrontendActionFactory(&refactor);
            auto result = Tool.run(Factory.get());
            if(result) {
                llvm::outs() << "Error occurs when doing code commit generation for file "
                             << Files[i] << ". \n";
            }
            resultSaveMutex.lock();
            resultMask |= result;
            resultSaveMutex.unlock();
        });
    }
    Pool.wait();
    if(0 == resultMask) {
        codeCommentInfoManager.doWriteOut();
    }
    return (resultMask == 0) ? 0 : 1;
}

int main(int argc, const char** argv)
{
    tooling::CommonOptionsParser OP(argc, argv, CodeCommentOptions);
    return Main_CodeCommentGen(OP);
}

